import os
import cv2
import numpy as np
from pathlib import Path
from tqdm import tqdm

def extract_and_save_depth_channel(input_path, output_path):
    """
    从四通道图像中提取深度通道并保存
    
    参数:
        input_path (str): 输入图像所在目录
        output_path (str): 输出深度图的目录
    """
    # 确保输出目录存在
    os.makedirs(output_path, exist_ok=True)
    
    # 获取所有图像文件
    image_extensions = ['.tiff', '.tif', '.png', '.jpg', '.jpeg']
    image_files = [
        f for f in os.listdir(input_path) 
        if any(f.lower().endswith(ext) for ext in image_extensions)
    ]
    
    if not image_files:
        print(f"错误: 在目录 {input_path} 中未找到图像文件")
        return
    
    print(f"找到 {len(image_files)} 个图像文件")
    
    # 遍历处理每个图像文件
    for image_file in tqdm(image_files, desc="处理进度"):
        file_path = os.path.join(input_path, image_file)
        file_base_name = Path(image_file).stem
        
        # 读取图像
        img = cv2.imread(file_path, cv2.IMREAD_UNCHANGED)
        
        if img is None:
            print(f"警告: 无法读取图像 {image_file}")
            continue
        
        # 检查是否为四通道图像
        if len(img.shape) != 3 or img.shape[2] != 4:
            print(f"警告: {image_file} 不是四通道图像 (shape: {img.shape})")
            continue
        
        # 提取深度通道（最后一个通道）
        depth_channel = img[:, :, 3]
        
        # 保存深度通道
        output_file = os.path.join(output_path, f"{file_base_name}_depth.png")
        cv2.imwrite(output_file, depth_channel)
    
    print(f"处理完成，深度图已保存到 {output_path}")

if __name__ == "__main__":
    # 设置参数
    INPUT_PATH = "/home/lhx/pcbsyn/data_zoo/splittruecolorwithdepthseg/images/train"  # 替换为输入图像路径
    OUTPUT_PATH = "/home/lhx/pcbsyn/data_zoo/splittruecolorwithdepthseg/imagevis"  # 替换为输出深度图路径
    
    # 执行提取
    extract_and_save_depth_channel(INPUT_PATH, OUTPUT_PATH)